from flask import Flask, send_file, jsonify, request, abort, send_from_directory
import os

app = Flask(__name__)

# 配置
IMAGE_FOLDER = 'images'  # 图片存放目录
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp'}  # 允许的图片格式
MAX_CONTENT_LENGTH = 16 * 1024 * 1024  # 最大上传文件大小 (16MB)
app.config['UPLOAD_FOLDER'] = IMAGE_FOLDER

# 确保图片目录存在
os.makedirs(IMAGE_FOLDER, exist_ok=True)

def allowed_file(filename):
    """检查文件扩展名是否合法"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/images', methods=['GET'])
def list_images():
    """获取图片列表"""
    try:
        images = []
        for filename in os.listdir(IMAGE_FOLDER):
            if allowed_file(filename):
                images.append({
                    'name': filename,
                    'url': f'/image/{filename}',
                    'size': os.path.getsize(os.path.join(IMAGE_FOLDER, filename))
                })
        return jsonify({
            'status': 'success',
            'count': len(images),
            'images': images
        })
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)}), 500

@app.route('/image/<path:filename>', methods=['GET'])
def get_image(filename):
    """获取原始图片"""
    # 防止路径遍历攻击
    if '..' in filename or filename.startswith('/'):
        abort(400, description="Invalid filename")
    
    if not allowed_file(filename):
        abort(400, description="Invalid file type")
    
    filepath = os.path.join(IMAGE_FOLDER, filename)
    if not os.path.exists(filepath):
        abort(404, description="Image not found")
    
    # 使用send_from_directory更安全
    return send_from_directory(IMAGE_FOLDER, filename)

@app.route('/upload', methods=['POST'])
def upload_image():
    """上传图片"""
    # 检查是否有文件被上传
    if 'file' not in request.files:
        return jsonify({'status': 'error', 'message': 'No file part'}), 400
    
    file = request.files['file']
    
    # 检查是否选择了文件
    if file.filename == '':
        return jsonify({'status': 'error', 'message': 'No selected file'}), 400
    
    # 验证文件类型
    if file and allowed_file(file.filename):
        # 安全处理文件名
        filename = secure_filename(file.filename)
        filepath = os.path.join(IMAGE_FOLDER, filename)
        
        # 避免覆盖已有文件
        counter = 1
        name, ext = os.path.splitext(filename)
        while os.path.exists(filepath):
            filename = f"{name}_{counter}{ext}"
            filepath = os.path.join(IMAGE_FOLDER, filename)
            counter += 1
        
        # 保存文件
        file.save(filepath)
        
        return jsonify({
            'status': 'success',
            'message': 'File uploaded successfully',
            'filename': filename,
            'url': f'/image/{filename}'
        }), 201
    
    return jsonify({'status': 'error', 'message': 'Invalid file type'}), 400

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=65301, debug=True)